# mnist_pinecone_knn.py

from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from collections import Counter
from tqdm import tqdm
import numpy as np
import logging
import time

# =================== 配置日志 ===================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()  # 输出到控制台
    ]
)

# =================== 初始化变量 ===================
API_KEY = "pcsk_7GvxrM_SVD5dH79WLzgGWd4JXQsdCoFHaXxsE3mJD83C4UXN8Ze1rRfkNC3ecQapfcd9DT"
INDEX_NAME = "mnist-knn-index"
DIMENSION = 64
METRIC = "euclidean"
CLOUD = "aws"
REGION = "us-east-1"

# =================== 加载 MNIST 数据 ===================
logging.info("开始加载 MNIST 手写数字数据集...")
digits = load_digits(n_class=10)
X, y = digits.data, digits.target

# 划分训练集（80%）和测试集（20%）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

logging.info(f"数据集划分完成：训练集 {len(X_train)} 条，测试集 {len(X_test)} 条")

# =================== 初始化 Pinecone ===================
pc = Pinecone(api_key=API_KEY)

# 删除已存在的索引（避免冲突）
if INDEX_NAME in pc.list_indexes().names():
    logging.info(f"索引 '{INDEX_NAME}' 已存在，正在删除...")
    pc.delete_index(INDEX_NAME)

# 创建新索引
logging.info(f"正在创建索引 '{INDEX_NAME}'...")
pc.create_index(
    name=INDEX_NAME,
    dimension=DIMENSION,
    metric=METRIC,
    spec=ServerlessSpec(cloud=CLOUD, region=REGION)
)
logging.info(f"索引 '{INDEX_NAME}' 创建成功")

# 连接索引
index = pc.Index(INDEX_NAME)
logging.info("正在连接到索引...")

# =================== 构建向量列表并上传 ===================
vectors = []
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 批量上传（带进度条）
BATCH_SIZE = 100
total_uploaded = 0

logging.info("开始上传训练数据（带进度条）...")

for i in tqdm(range(0, len(vectors), BATCH_SIZE), desc="📤 上传向量"):
    batch = vectors[i:i + BATCH_SIZE]
    index.upsert(vectors=batch)
    total_uploaded += len(batch)

logging.info(f"✅ 成功创建索引，并上传了 {total_uploaded} 条数据")

# =================== 测试阶段：KNN 查询（k=11）===================
K = 11
correct_predictions = 0

logging.info(f"开始测试 KNN 分类性能（k={K}），使用 {len(X_test)} 个测试样本...")

for i in tqdm(range(len(X_test)), desc="🔍 查询相似向量"):
    query_vector = X_test[i].tolist()
    
    # 查询最相似的 k 个向量
    try:
        result = index.query(
            vector=query_vector,
            top_k=K,
            include_metadata=True
        )
        
        # 提取标签（跳过自身可能匹配）
        labels = []
        for match in result['matches']:
            if 'metadata' in match and 'label' in match['metadata']:
                labels.append(match['metadata']['label'])
        
        # 如果没找到标签，跳过
        if not labels:
            continue
        
        # 投票决定预测结果（多数投票）
        predicted_label = Counter(labels).most_common(1)[0][0]
        
        # 比较真实标签
        if predicted_label == y_test[i]:
            correct_predictions += 1
            
    except Exception as e:
        logging.warning(f"查询第 {i} 个样本时出错: {e}")

# =================== 计算准确率 ===================
accuracy = correct_predictions / len(X_test)
logging.info(f"✅ 当 k={K} 时，使用 Pinecone 的准确率为: {accuracy:.4f} ({correct_predictions}/{len(X_test)})")

# =================== 可选：删除索引（释放资源）===================
# pc.delete_index(INDEX_NAME)
# logging.info(f"🗑️ 索引 '{INDEX_NAME}' 已删除")